package com.atguigu.easyexcel;

import java.util.ArrayList;
import java.util.Random;

public class KMeans {
    private int k;  // 簇的个数
    private int n;  // 样本数
    private int m;  // 样本维度
    private int maxIter;  // 最大迭代次数
    private double[][] data;  // 样本数据
    private int[] labels;  // 样本标签
    private double[][] centers;  // 簇中心
    private int[] sizes;  // 簇大小
    private int capacity;  // 宿舍容量

    public KMeans(int k, int n, int m, int maxIter, double[][] data, int capacity) {
        this.k = k;
        this.n = n;
        this.m = m;
        this.maxIter = maxIter;
        this.data = data;
        this.labels = new int[n];
        this.centers = new double[k][m];
        this.sizes = new int[k];
        this.capacity = capacity;
    }

    // 初始化簇中心
    private void initCenters() {
        Random random = new Random();
        ArrayList<Integer> taken = new ArrayList<Integer>();
        for (int i = 0; i < k; i++) {
            int idx;
            do {
                idx = random.nextInt(n);
            } while (taken.contains(idx));
            taken.add(idx);
            for (int j = 0; j < m; j++) {
                centers[i][j] = data[idx][j];
            }
        }
    }

    // 计算样本到簇中心的距离
    private double dist(double[] a, double[] b) {
        double sum = 0;
        for (int i = 0; i < m; i++) {
            sum += Math.pow(a[i] - b[i], 2);
        }
        return Math.sqrt(sum);
    }

    // 聚类
    private void cluster() {
        for (int i = 0; i < n; i++) {
            double minDist = Double.MAX_VALUE;
            int minIdx = -1;
            for (int j = 0; j < k; j++) {
                double d = dist(data[i], centers[j]);
                if (d < minDist) {
                    minDist = d;
                    minIdx = j;
                }
            }
            int idx = minIdx;
            while (sizes[idx] >= capacity) {
                minDist = Double.MAX_VALUE;
                minIdx = -1;
                for (int j = 0; j < k; j++) {
                    double d = dist(data[i], centers[j]);
                    if (d < minDist && sizes[j] < capacity) {
                        minDist = d;
                        minIdx = j;
                    }
                }
                if (minIdx < 0) {
                    int remaining = n - capacity * k;
                    int increment = (remaining + sizes[idx]) / (k - idx);
                    for (int j = idx; j < k; j++) {
                        sizes[j] += increment;
                    }
                    break;
                }
                idx = minIdx;
            }
            labels[i] = idx;
            sizes[idx]++;
        }
    }

    // 更新簇中心
    private void updateCenters() {
        for (int i = 0; i < k; i++) {
            for (int j = 0; j < m; j++) {
                centers[i][j] = 0;
            }
        }
        for (int i = 0; i < n; i++) {
            int idx = labels[i];
            for (int j = 0; j < m; j++) {
                centers[idx][j] += data[i][j];
            }
        }
        for (int i = 0; i < k; i++) {
            if (sizes[i] > 0) {
                for (int j = 0; j < m; j++) {
                    centers[i][j] /= sizes[i];
                }
            }
        }
    }

    // 迭代
    public void run() {
        initCenters();
        for (int iter = 0; iter < maxIter; iter++) {
            sizes = new int[k];
            cluster();
            updateCenters();
        }
    }

    // 获取簇标签
    public int[] getLabels() {
        return labels;
    }

    public static void main(String[] args) {
        int capacity = 6;//宿舍容量
        int n = 100; // 样本数
        int k = n % capacity ==0 ?n / capacity:n / capacity + 1; // 簇的个数
        int m = 8; // 样本维度
        int maxIter = 100; // 最大迭代次数
        double[][] data = new double[n][m]; // 样本数据
        Random random = new Random();
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < m; j++) {
                data[i][j] = random.nextDouble();
            }
        }

        KMeans kMeans = new KMeans(k, n, m, maxIter, data, capacity);
        kMeans.run();
        int[][] clusters = new int[k][n]; // 二维数组存储聚类结果
        for (int i = 0; i < n; i++) {
            int clusterIndex = kMeans.getLabels()[i]; // 获取该样本所属的簇编号
            clusters[clusterIndex][i] = 1; // 将该样本的聚类结果存储到对应的簇中
        }
        for (int i = 0; i < k; i++) {
            System.out.print("簇 " + i + " 中的样本：");
            for (int j = 0; j < n; j++) {
                if (clusters[i][j] == 1) {
                    System.out.print(j + " ");
                }
            }
            System.out.println();
        }
    }
}